from datasets import load_dataset
import pandas as pd
from collections import defaultdict
import re
import os
from openai import OpenAI
import time

def count_chinese_chars(text):
    """统计中文字符数"""
    return len(re.findall(r'[\u4e00-\u9fff]', text))

def count_english_words(text):
    """统计英文单词数"""
    return len(text.split())

def save_to_excel(data, filename, max_rows=1000000):
    """将数据分批保存到多个Excel文件"""
    if not data:
        return
    
    # 创建输出目录
    os.makedirs('output3', exist_ok=True)
    
    # 计算需要多少个文件
    total_chunks = (len(data) + max_rows - 1) // max_rows
    
    for i in range(total_chunks):
        start_idx = i * max_rows
        end_idx = min((i + 1) * max_rows, len(data))
        chunk_data = data[start_idx:end_idx]
        
        # 生成文件名
        if total_chunks > 1:
            chunk_filename = f'output3/{filename}_{i+1}.xlsx'
        else:
            chunk_filename = f'output3/{filename}.xlsx'
        
        # 保存数据
        df = pd.DataFrame(chunk_data)
        df.to_excel(chunk_filename, index=False)
        print(f"保存数据到 {chunk_filename}, 共 {len(chunk_data)} 条")

def classify_style_and_domain(text):
    """根据数据集来源直接分类风格和领域"""
    # TED演讲和Flores中的教育内容
    return '教育翻译', '请将以下教育相关的中文内容翻译成英文：'

def generate_education_content(client, num_samples=10):
    """使用 GPT-4 生成教育相关的中文内容和英文翻译"""
    translation_data = []
    
    for i in range(num_samples):
        try:
            # 生成中文内容
            zh_response = client.chat.completions.create(
                model="gpt-4",
                messages=[
                    {"role": "system", "content": "你是一个教育内容创作者。"},
                    {"role": "user", "content": "请生成一段300字左右的关于教育主题的中文段落。主题可以包括教学方法、学习理论、教育创新等。"}
                ]
            )
            zh_text = zh_response.choices[0].message.content

            # 生成英文翻译
            en_response = client.chat.completions.create(
                model="gpt-4",
                messages=[
                    {"role": "system", "content": "你是一个专业的中英翻译。"},
                    {"role": "user", "content": f"请将以下中文内容翻译成英文：\n\n{zh_text}"}
                ]
            )
            en_text = en_response.choices[0].message.content

            # 添加到数据集
            translation_data.append({
                '指令': '请将以下教育相关的中文内容翻译成英文：',
                '原文': zh_text,
                '译文': en_text
            })

            print(f"已生成第 {i+1} 条数据")
            time.sleep(1)  # 避免触发 API 限制

        except Exception as e:
            print(f"生成第 {i+1} 条数据时出错: {str(e)}")
            continue

    return translation_data

def main():
    # 初始化 OpenAI 客户端
    client = OpenAI(api_key='your-api-key-here')  # 替换为你的 API key
    
    # 生成数据
    print("开始生成教育相关内容...")
    translation_data = generate_education_content(client, num_samples=10)  # 生成10条数据
    
    # 保存数据
    if translation_data:
        save_to_excel(translation_data, 'ai_generated_education_dataset')
        print("\n所有数据已保存到 output3 目录")
    else:
        print("未生成任何数据")

if __name__ == "__main__":
    main()

# 加载多个数据集
print("正在加载数据集...")
datasets = {
    'ted_ed': load_dataset("iwslt2017", "iwslt2017-zh-en", trust_remote_code=True),  # 添加 trust_remote_code 参数
    # 'amazon': load_dataset("amazon_reviews_multi", "zh",trust_remote_code=True),  # 营销相关评论
    # 'jd': load_dataset("DAMO-NLP-SG/jd_product_comment",trust_remote_code=True),  # 营销相关评论
    # 'flores': load_dataset("facebook/flores", "zho_Hans-eng")  # 包含多领域的翻译数据
}

# 按不同类型收集数据
translation_data = defaultdict(list)
total_processed = 0
length_filtered = 0

# 字数限制
MIN_ZH_LENGTH = 300
MAX_ZH_LENGTH = 2000
MIN_EN_LENGTH = 20
MAX_EN_LENGTH = 3000

# 处理每个数据集
for dataset_name, ds in datasets.items():
    print(f"\n处理 {dataset_name} 数据集...")
    
    if dataset_name == 'ted_ed':
        data = ds['train']
        for example in data:
            total_processed += 1
            if total_processed % 10000 == 0:
                print(f"已处理 {total_processed} 条数据...")
            
            zh_text = example['translation']['zh']
            en_text = example['translation']['en']
            
            zh_length = count_chinese_chars(zh_text)
            en_length = count_english_words(en_text)
            
            if MIN_ZH_LENGTH <= zh_length <= MAX_ZH_LENGTH and MIN_EN_LENGTH <= en_length <= MAX_EN_LENGTH:
                length_filtered += 1
                category, instruction = classify_style_and_domain(en_text)
                if category is not None:
                    translation_data[category].append({
                        '指令': instruction,
                        '原文': zh_text,
                        '译文': en_text
                    })

    elif dataset_name == 'flores':
        data = ds['dev']  # flores 使用 dev 集
        for example in data:
            total_processed += 1
            if total_processed % 10000 == 0:
                print(f"已处理 {total_processed} 条数据...")
            
            zh_text = example['zho_Hans']
            en_text = example['eng']
            
            zh_length = count_chinese_chars(zh_text)
            en_length = count_english_words(en_text)
            
            if MIN_ZH_LENGTH <= zh_length <= MAX_ZH_LENGTH and MIN_EN_LENGTH <= en_length <= MAX_EN_LENGTH:
                length_filtered += 1
                category, instruction = classify_style_and_domain(en_text)
                if category is not None:
                    translation_data[category].append({
                        '指令': instruction,
                        '原文': zh_text,
                        '译文': en_text
                    })

    elif dataset_name == 'amazon':
        data = ds['train']
        for example in data:
            total_processed += 1
            if total_processed % 10000 == 0:
                print(f"已处理 {total_processed} 条数据...")
            
            zh_text = example['review_body']
            
            # Amazon数据集只有中文，我们只收集中文内容
            zh_length = count_chinese_chars(zh_text)
            
            if MIN_ZH_LENGTH <= zh_length <= MAX_ZH_LENGTH:
                length_filtered += 1
                translation_data['营销翻译'].append({
                    '指令': '请将以下营销相关的中文内容翻译成英文：',
                    '原文': zh_text,
                    '译文': ''  # 留空，因为没有对应的英文翻译
                })
    
    elif dataset_name == 'jd':
        data = ds['train']
        for example in data:
            total_processed += 1
            if total_processed % 10000 == 0:
                print(f"已处理 {total_processed} 条数据...")
            
            zh_text = example['comment']  # JD数据集的评论字段
            
            # JD数据集只有中文，我们只收集中文内容
            zh_length = count_chinese_chars(zh_text)
            
            if MIN_ZH_LENGTH <= zh_length <= MAX_ZH_LENGTH:
                length_filtered += 1
                translation_data['营销翻译'].append({
                    '指令': '请将以下营销相关的中文内容翻译成英文：',
                    '原文': zh_text,
                    '译文': ''  # 留空，因为没有对应的英文翻译
                })

print(f"\n符合字数要求的数据: {length_filtered} 条")
print("\n各类别数据统计:")
for category, data in translation_data.items():
    print(f"{category}: {len(data)}条")

# 分别保存每个类别的数据
for category, data in sorted(translation_data.items()):
    if len(data) > 0:
        save_to_excel(data, f'translation_dataset_{category}')

print("\n所有数据已保存到 output3 目录")